#include "encmngr.h"
#include <aes.h>
#include <algorithm>
#include "incognito.h"

const unsigned int max_size=1024; //Maximum size of encrypted data (WARNING: it must be divisible by 16) FIXME

// {{{ Calculates SHA-1 of Diffie-Hellman Key concatended with sequence number
void calc_aes_key(char *key,bigint zz, unsigned int seq)
{
	sha1ctx sc;
	sc.update(zz.getraw(),zz.getraw().len());
	sc.update(&seq,sizeof(seq));
	bzero(key, sha1::hashsize);
	sc.final(key);
}
// }}}

// {{{ Encryption
// {{{ void encmngr::encrypt(chord_node target, ptr<incognito_raw_request> raw_msg, cb_encrypt cb)
void encmngr::encrypt(chord_node target, ptr<incognito_raw_request> raw_msg, cb_encrypt cb)
{
	warnx << "encmngr::encrypt();\n";
	return _keymngr->get_key(_ic,target,wrap(mkref(this),&encmngr::encrypt_get_key_cb,target.x,raw_msg,cb));
}
// {{{ void encmngr::encrypt(chordID target, ptr<incognito_raw_request> raw_msg, cb_encrypt cb)
void encmngr::encrypt(chordID target, ptr<incognito_raw_request> raw_msg, cb_encrypt cb)
{
	warnx << "encmngr::encrypt_no_new_key();\n";
	return _keymngr->get_key(_ic,target,wrap(mkref(this),&encmngr::encrypt_get_key_cb,target,raw_msg,cb));
}
// }}}

// {{{ void encmngr::encrypt_get_key_cb(chordID target, ptr<incognito_raw_request> raw_msg, cb_encrypt cb, bigint k, getkeyres res)
void encmngr::encrypt_get_key_cb(chordID target, ptr<incognito_raw_request> raw_msg, cb_encrypt cb, bigint k, getkeyres res)
{
	ptr<incognito_encrypted_request> result=0;
	warnx << "encmngr::encrypt_get_key_cb(" << target << "\n\t" << k << ");\n";
	
	if (res!=INCOGNITO_KEY_OK) {				// If there is an error...
		return cb(result, INCOGNITO_EM_ENCRYPT_NO_KEY);		
	}

	if (_encmap.find(target) == _encmap.end())		// Add new enctyption parameters if not exists
		_encmap[target]=encparam();

	unsigned int seq=_encmap[target].send_seq++;		// Calculate sequence number

	char h[sha1::hashsize];					// Calculate encryption key
	calc_aes_key(h,k, seq);

#if 0
	bigint nn;
	mpz_set_rawmag_be(&nn, h, sha1::hashsize);  // For big endian
	warnx << "sha1=" << nn << "\n";
#endif

	// XDR - start
	XDR xdrs;
	char buf[10240];						// FIXME: Constant buffer size...
	xdrmem_create(&xdrs,buf,10240,XDR_ENCODE);			// FIXME: ^^^^^^^^^^^^^^^^^^^^
	xdr_incognito_raw_request(&xdrs,raw_msg);			// Store request

	if (xdr_getpos(&xdrs) > max_size) {				// Check if it's not too big
		warnx << "encmngr::encrypt_get_key_cb("
			<< target << ") - xdr_getpos(&xdrs) > max_size FIXME: "
			<< __FILE__ << ":" << __LINE__ << "\n";
		return cb(result, INCOGNITO_EM_ENCRYPT_REQUEST_TOO_BIG);
	}

	while (xdr_getpos(&xdrs) < max_size) {				// Fill free space with random data
		unsigned int x=random();
		xdr_int(&xdrs,&x);
	}
	xdr_destroy(&xdrs);
	// XDR - stop

	// Encryption - start
	aes_e a;
	a.setkey(h,16);

	for(unsigned int i=0;i<max_size;i+=16)				// Encrypt buffer
		a.encipher_bytes(buf+i);
	// Encryption - stop

        result = New refcounted<incognito_encrypted_request> ();	// Prepare response
        result->seq_num=seq; 						// caller_id must be filled later
        result->enc_data.setsize(max_size);
        memcpy (result->enc_data.base (), buf, max_size);


	return cb(result,INCOGNITO_EM_ENCRYPT_OK);			// Call callback
}
// }}} 
// }}}
// {{{ Decryption
void encmngr::decrypt(ptr<incognito_encrypted_request> enc_msg, cb_decrypt cb)
{
	warnx << "encmngr::decrypt();\n";
	_keymngr->get_key(_ic,enc_msg->caller_ID,wrap(mkref(this),&encmngr::decrypt2_get_key_cb,enc_msg,cb));

}
void encmngr::decrypt2_get_key_cb(ptr<incognito_encrypted_request> enc_msg, cb_decrypt cb, bigint k, getkeyres res)
{
	warnx << "encmngr::decrypt_get_key_cb()\n";

	ptr<incognito_raw_request> raw_msg=0;
	if (res!=INCOGNITO_KEY_OK) {					// If there is an error...
		cb(raw_msg,INCOGNITO_EM_DECRYPT_NO_KEY);
		return;
	}

	char h[sha1::hashsize];						// Calculate encryption key
	calc_aes_key(h,k,enc_msg->seq_num);
	
	char buf[enc_msg->enc_data.size()];				// Copy encrypted data from message
	memcpy(buf,enc_msg->enc_data.base(),enc_msg->enc_data.size());	//  to buffer

	// Decryption - start
	aes dec;
	dec.setkey(h,16);

	for(unsigned int i=0;i<enc_msg->enc_data.size();i+=16)		// Decrypt buffer
		dec.decipher_bytes(buf+i);
	// Decryption - stop
	
	// XDR - start
	XDR xdrs;
	xdrmem_create(&xdrs,buf,enc_msg->enc_data.size(),XDR_DECODE);	// Create XDR decoding stream

        raw_msg = New refcounted<incognito_raw_request> ();		// Allocate memory for response

	if (!xdr_incognito_raw_request(&xdrs,raw_msg))			// Extract data from XDR stream
		return cb(raw_msg,INCOGNITO_EM_DECRYPT_XDR_FAILED);	// Return error on fail 

	xdr_destroy(&xdrs);
	// XDR - stop

	cb(raw_msg,INCOGNITO_EM_DECRYPT_OK);				// Return decrypted data
}
// }}}


//  {{{ Obsolete methods
void encmngr::encrypt_and_send_store(chord_node &n,ptr<incognito_store_arg> req)
{
	warnx << "encmngr::encrypt_and_send_store();\n";
	_keymngr->get_key(_ic,n,wrap(mkref(this),&encmngr::encrypt_and_send_store_get_key_cb,n,req));
}
void encmngr::encrypt_and_send_store_get_key_cb(chord_node n,ptr<incognito_store_arg> req,bigint k,getkeyres res)
{
	warnx << "encmngr::encrypt_and_send_store_get_key_cb(" << n.x << "\n\t" << k << ");\n";
	if (res!=INCOGNITO_KEY_OK) {
		warnx << "encmngr::encrypt_and_send_store_get_key_cb("
			<< n.x << ") - res!=INCOGNITO_KEY_OK FIXME: "
			<< __FILE__ << ":" << __LINE__ << "\n";
		return; // FIXME
	}
	if (_encmap.find(n.x) == _encmap.end())
		_encmap[n.x]=encparam();
	warnx << "zz=" << k << "\n";
	warnx << "strlen(zz.cstr())=" << strlen(k.cstr()) << "\n";
	warnx << "zz.getraw().len()=" << k.getraw().len() << "\n";
	warnx << "zz.nbits()=" << k.nbits() << "\n";

	char h[sha1::hashsize];
	unsigned int seq=_encmap[n.x].send_seq++;
	calc_aes_key(h,k, seq);
	
	bigint nn;
	mpz_set_rawmag_be(&nn, h, sha1::hashsize);  // For big endian
	warnx << "sha1=" << nn << "\n";

	//XDR - start
	XDR xdrs;
	char buf[10240];
	xdrmem_create(&xdrs,buf,10240,XDR_ENCODE);
	warnx << "xdr_getpos(&xdrs)=" << xdr_getpos(&xdrs) << "\n";
	warnx << "xdr_incognito_store_arg(&xdrs,sarg)=" << xdr_incognito_store_arg(&xdrs,req) << "\n";
	warnx << "xdr_getpos(&xdrs)=" << xdr_getpos(&xdrs) << " 0\n";
	if (xdr_getpos(&xdrs) > max_size) {
		warnx << "encmngr::encrypt_and_send_store_get_key_cb("
			<< n.x << ") - xdr_getpos(&xdrs) > max_size FIXME: "
			<< __FILE__ << ":" << __LINE__ << "\n";
		return;	// FIXME
	}
	while (xdr_getpos(&xdrs) < max_size) {
		unsigned int x=random();
		xdr_int(&xdrs,&x);
	}
	warnx << "xdr_getpos(&xdrs)=" << xdr_getpos(&xdrs) << " 1\n";
	xdr_destroy(&xdrs);

	aes_e a;
	a.setkey(h,16);

	char tab[]={
		'0','1','2','3',
		'4','5','6','7',
		'8','9','A','B',
		'C','D','E','F'};

	warn << "Raw:\t";
	for (unsigned int i=0;i<min(max_size,(unsigned int)64);++i) {
		char p[3]={tab[((unsigned char)buf[i])>>4],tab[((unsigned char)buf[i])&0x0f],0};
		warnx << p;
	}
	warnx << "\n";

	for(unsigned int i=0;i<max_size;i+=16)
		a.encipher_bytes(buf+i);
	
	_ic->store_encrypted_data(n,seq,buf,max_size);
	
	warn << "Encrypted:\t";
	for (unsigned int i=0;i<min(max_size,(unsigned int)64);++i) {
		char p[3]={tab[((unsigned char)buf[i])>>4],tab[((unsigned char)buf[i])&0x0f],0};
		warnx << p;
	}
	warnx << "\n";

	aes dec;
	dec.setkey(h,16);

	for(unsigned int i=0;i<max_size;i+=16)
		dec.decipher_bytes(buf+i);
	
	warn << "Decrypted:\t";
	for (unsigned int i=0;i<min(max_size,(unsigned int)64);++i) {
		char p[3]={tab[((unsigned char)buf[i])>>4],tab[((unsigned char)buf[i])&0x0f],0};
		warnx << p;
	}
	warnx << "\n";
	
}

void encmngr::decrypt(ptr<incognito_enc_store_arg> arg, cbdecrypt cb)
{
	warnx << "encmngr::decrypt();\n";
	_keymngr->get_key(_ic,arg->caller_ID,wrap(mkref(this),&encmngr::decrypt_get_key_cb,arg,cb));
}

void encmngr::decrypt_get_key_cb(ptr<incognito_enc_store_arg> arg, cbdecrypt cb, bigint k, getkeyres res)
{
	incognito_store_arg p;
	warnx << "encmngr::decrypt_get_key_cb()\n";
	if (res!=INCOGNITO_KEY_OK) {
		warnx << "encmngr::decrypt_get_key_cb("
			<< arg->caller_ID << ") - res!=INCOGNITO_KEY_OK FIXME: "
			<< __FILE__ << ":" << __LINE__ << "\n";
		cb(p,INCOGNITO_DECRYPT_NO_KEY);
		return;
	}
	char h[sha1::hashsize];
	calc_aes_key(h,k,arg->seq_num);
	
	char buf[arg->enc_data.size()];
	memcpy(buf,arg->enc_data.base(),arg->enc_data.size());

	aes dec;
	dec.setkey(h,16);

	for(unsigned int i=0;i<arg->enc_data.size();i+=16)
		dec.decipher_bytes(buf+i);
	
	char tab[]={
		'0','1','2','3',
		'4','5','6','7',
		'8','9','A','B',
		'C','D','E','F'};
	warn << "encmngr::decrypt_get_key_cb() - Decrypted:\n";
	for (unsigned int i=0;i<max_size;++i) {
		char p[3]={tab[((unsigned char)buf[i])>>4],tab[((unsigned char)buf[i])&0x0f],0};
		warnx << p;
	}
	warnx << "\n";
	
	XDR xdrs;
	xdrmem_create(&xdrs,buf,arg->enc_data.size(),XDR_DECODE);
	if (!xdr_incognito_store_arg(&xdrs,&p)) {
		warnx << "encmngr::decrypt_get_key_cb("
			<< arg->caller_ID << ") - !xdr_incognito_store_arg(&xdrs,p) FIXME: "
			<< __FILE__ << ":" << __LINE__ << "\n";
		cb(p,INCOGNITO_DECRYPT_XDR_FAILED);
		return;

	}
	xdr_destroy(&xdrs);
	cb(p,INCOGNITO_DECRYPT_OK);
}
// }}}

void encmngr::print_stats()
{
	warnx << "encmngr::print_stats() - start -------------------\n";
	warnx << "encmngr::print_stats() - stop  -------------------\n";
}
// vim: foldmethod=marker
